package cc.git.liuyan.customeraiagent.core.embeddingmodel;

import java.util.ArrayList;
import java.util.List;

public class SplitUtil {
    public static List<String> splitText(String text) {
        return splitText(text, 400, 40);
    }

    public static List<String> splitText(String text, int chunkSize, int overLapSize) {
        List<String> chunks = new ArrayList<>();
        int textLength = text.length();
        String lastChunk = null;
        for (int i = 0; i < textLength; i += chunkSize) {
            int end = Math.min(i + chunkSize, textLength);
            String chunk = text.substring(i, end);
            if (lastChunk != null) {
                int overLapLength = Math.min(lastChunk.length(), overLapSize);
                chunks.add(lastChunk.substring(lastChunk.length() - overLapLength) + chunk);
            } else {
                chunks.add(chunk);
            }
            lastChunk = chunk;
        }
        return chunks;
    }

    public static List<EmbeddingModelInputData> splitStringInputData(EmbeddingModelInputData inputData) {
        return splitStringInputData(inputData, 400, 40);
    }

    public static List<EmbeddingModelInputData> splitStringInputData(EmbeddingModelInputData inputData, int chunkSize, int overLapSize) {
        List<EmbeddingModelInputData> chunks = new ArrayList<>();
        int textLength = inputData.getInputData().length();
        String lastChunk = null;
        for (int i = 0; i < textLength; i += chunkSize) {
            int end = Math.min(i + chunkSize, textLength);
            String chunk = inputData.getInputData().substring(i, end);
            if (lastChunk != null) {
                int overLapLength = Math.min(lastChunk.length(), overLapSize);
                chunks.add(new EmbeddingModelInputData(lastChunk.substring(lastChunk.length() - overLapLength) + chunk));
            } else {
                chunks.add(new EmbeddingModelInputData(chunk));
            }
            lastChunk = chunk;
        }
        return chunks;
    }

}
